import sys
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
from matplotlib.widgets import Button, Slider
import akshare as ak
import datetime
from datetime import date
from dateutil.relativedelta import relativedelta  # 正确导入方式
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.preprocessing import MinMaxScaler, StandardScaler
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import ParameterGrid
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
                            QLabel, QComboBox, QSpinBox, QPushButton, QTextEdit, QProgressBar,
                            QMessageBox, QTabWidget, QGridLayout, QCheckBox, QGroupBox,
                            QDoubleSpinBox, QSlider, QFileDialog, QSplitter, QToolBar, QAction)
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from typing import List, Dict, Tuple, Optional
import requests  # 新增：用于请求节假日API

import sys
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
import akshare as ak
import datetime
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.preprocessing import MinMaxScaler, StandardScaler
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import ParameterGrid
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
                            QLabel, QComboBox, QSpinBox, QPushButton, QTextEdit, QProgressBar,
                            QMessageBox, QTabWidget, QGridLayout, QCheckBox, QGroupBox,
                            QDoubleSpinBox, QSlider, QFileDialog, QSplitter)
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from typing import List, Dict, Tuple, Optional

from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT



import sys
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.figure import Figure
from matplotlib.widgets import Button, Slider  # 导入缩放相关组件
import akshare as ak
import datetime
import os
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.preprocessing import MinMaxScaler, StandardScaler
from sklearn.metrics import mean_absolute_error, mean_squared_error, r2_score
from sklearn.model_selection import ParameterGrid
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, QHBoxLayout, 
                            QLabel, QComboBox, QSpinBox, QPushButton, QTextEdit, QProgressBar,
                            QMessageBox, QTabWidget, QGridLayout, QCheckBox, QGroupBox,
                            QDoubleSpinBox, QFileDialog, QSplitter, QToolBar, QAction)
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from typing import List, Dict, Tuple, Optional


# 设置中文字体
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]
plt.rcParams["axes.unicode_minus"] = False  # 解决负号显示问题

class AShareTradingDayChecker:
    """中国A股交易日检查工具"""
    def __init__(self):
        self.cache = {}
        self.holiday_api = "http://tool.bitefu.net/jiari/"
    
    def is_trading_day(self, date: date) -> bool:
        """判断指定日期是否为A股交易日（排除周末和节假日）"""
        if date.weekday() >= 5:  # 排除周末
            return False
        return self._is_working_day(date)
    
    def _is_working_day(self, date: date) -> bool:
        """检查是否为工作日（使用API或备用规则）"""
        date_str = date.strftime("%Y%m%d")
        
        # 优先使用缓存
        if date_str in self.cache:
            return self.cache[date_str]
        
        try:
            # 请求节假日API
            response = requests.get(self.holiday_api, params={"d": date_str})
            if response.status_code == 200:
                result = int(response.text)
                # 0: 工作日, 1: 休息日, 2: 节假日（均非交易日）
                is_working = result == 0
                self.cache[date_str] = is_working
                return is_working
        except Exception as e:
            print(f"节假日API请求失败: {e}")
        
        # API不可用时的备用规则（覆盖常见节假日，需根据实际情况更新）
        return self._fallback_is_working_day(date)
    
    def _fallback_is_working_day(self, date: date) -> bool:
        """备用节假日判断（仅示例，需根据最新放假安排调整）"""
        # 定义常见节假日（月, 日）- 2025年示例，实际应动态更新
        holidays = [
            (1, 1),   # 元旦
            (2, 1), (2, 2), (2, 3), (2, 4), (2, 5), (2, 6), (2, 7),  # 春节
            (4, 5),   # 清明
            (5, 1),   # 劳动节
            (6, 10), (6, 11), (6, 12),  # 端午
            (9, 29), (9, 30), (10, 1), (10, 2), (10, 3), (10, 4), (10, 5)  # 中秋国庆
        ]
        month_day = (date.month, date.day)
        return month_day not in holidays
    
    def get_trading_days(self, start_date: date, end_date: date) -> list[date]:
        """获取指定区间内的所有交易日"""
        trading_days = []
        current_date = start_date
        while current_date <= end_date:
            if self.is_trading_day(current_date):
                trading_days.append(current_date)
            current_date += relativedelta(days=1)
        return trading_days

class Attention(nn.Module):
    """注意力机制层"""
    def __init__(self, hidden_size):
        super(Attention, self).__init__()
        self.hidden_size = hidden_size
        self.attn = nn.Linear(self.hidden_size * 2, hidden_size)
        self.v = nn.Parameter(torch.rand(hidden_size))
        stdv = 1. / torch.sqrt(self.v.size(0))
        self.v.data.uniform_(-stdv, stdv)
        
    def forward(self, hidden, encoder_outputs):
        timestep = encoder_outputs.size(0)
        hidden = hidden.repeat(timestep, 1, 1).transpose(0, 1)
        encoder_outputs = encoder_outputs.transpose(0, 1)
        attn_energies = self.score(hidden, encoder_outputs)
        return torch.softmax(attn_energies, dim=1).unsqueeze(1)
    
    def score(self, hidden, encoder_outputs):
        energy = torch.tanh(self.attn(torch.cat([hidden, encoder_outputs], 2)))
        energy = energy.transpose(1, 2)
        v = self.v.repeat(encoder_outputs.size(0), 1).unsqueeze(1)
        energy = torch.bmm(v, energy)
        return energy.squeeze(1)

class LSTMModel(nn.Module):
    """改进的LSTM模型，支持注意力机制和多层结构"""
    def __init__(self, input_size=1, hidden_size=50, num_layers=2, output_size=1, use_attention=False):
        super(LSTMModel, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.use_attention = use_attention
        
        # LSTM层
        self.lstm = nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=0.2)
        
        # 注意力机制
        if use_attention:
            self.attention = Attention(hidden_size)
            
        # 全连接层
        self.fc = nn.Linear(hidden_size, output_size)
        
    def forward(self, x):
        # 初始化隐藏状态
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        c0 = torch.zeros(self.num_layers, x.size(0), self.hidden_size).to(x.device)
        
        # LSTM前向传播
        lstm_out, _ = self.lstm(x, (h0, c0))  # [batch, seq_len, hidden_size]
        
        # 应用注意力机制
        if self.use_attention:
            attn_weights = self.attention(lstm_out[:, -1:, :], lstm_out)
            context = attn_weights.bmm(lstm_out)  # [batch, 1, hidden_size]
            output = self.fc(context.squeeze(1))
        else:
            # 只取最后一个时间步的输出
            output = self.fc(lstm_out[:, -1, :])
            
        return output

class HyperparameterTuningThread(QThread):
    """超参数调优线程"""
    progress_updated = pyqtSignal(int)
    message_updated = pyqtSignal(str)
    tuning_completed = pyqtSignal(dict)
    
    def __init__(self, stock_predictor, X_train, y_train, X_test, y_test):
        super().__init__()
        self.stock_predictor = stock_predictor
        self.X_train = X_train
        self.y_train = y_train
        self.X_test = X_test
        self.y_test = y_test
        
        # 定义超参数搜索空间
        self.param_grid = {
            'hidden_size': [50, 100],
            'num_layers': [2, 3],
            'epochs': [30, 50],
            'batch_size': [16, 32],
            'learning_rate': [0.001, 0.0005]
        }
        
    def run(self):
        try:
            total_trials = len(ParameterGrid(self.param_grid))
            best_rmse = float('inf')
            best_params = None
            results = []
            
            self.message_updated.emit(f"开始超参数调优，共{total_trials}种组合...")
            
            for i, params in enumerate(ParameterGrid(self.param_grid)):
                self.message_updated.emit(f"尝试参数组合 {i+1}/{total_trials}: {params}")
                
                # 构建模型
                self.stock_predictor.build_model(
                    hidden_size=params['hidden_size'],
                    num_layers=params['num_layers']
                )
                
                # 训练模型
                self.stock_predictor.train_model(
                    self.X_train, self.y_train,
                    epochs=params['epochs'],
                    batch_size=params['batch_size'],
                    learning_rate=params['learning_rate']
                )
                
                # 评估模型
                metrics = self.stock_predictor.evaluate_model(self.X_test, self.y_test)
                results.append({**params, **metrics})
                
                # 更新最佳参数
                if metrics['RMSE'] < best_rmse:
                    best_rmse = metrics['RMSE']
                    best_params = params
                    
                self.progress_updated.emit(int((i+1) / total_trials * 100))
            
            self.message_updated.emit(f"超参数调优完成！最佳参数: {best_params}, RMSE: {best_rmse:.4f}")
            self.tuning_completed.emit({
                'best_params': best_params,
                'results': results
            })
            
        except Exception as e:
            self.message_updated.emit(f"调优过程出错: {str(e)}")

class PredictionThread(QThread):
    """后台运行预测任务的线程"""
    progress_updated = pyqtSignal(int)
    message_updated = pyqtSignal(str)
    prediction_completed = pyqtSignal(object, object, object, object)
    
    def __init__(self, stock_predictor, epochs, batch_size, sequence_length, learning_rate, use_attention, use_technical_indicators, use_standard_scaler):
        super().__init__()
        self.stock_predictor = stock_predictor
        self.epochs = epochs
        self.batch_size = batch_size
        self.sequence_length = sequence_length
        self.learning_rate = learning_rate
        self.use_attention = use_attention
        self.use_technical_indicators = use_technical_indicators
        self.use_standard_scaler = use_standard_scaler
        
    def run(self):
        try:
            self.message_updated.emit("开始获取数据...")
            self.stock_predictor.fetch_data(use_technical_indicators=self.use_technical_indicators)
            self.progress_updated.emit(5)
            
            self.message_updated.emit("开始数据预处理...")
            scaler_type = 'standard' if self.use_standard_scaler else 'minmax'
            X_train, y_train, X_test, y_test, scaled_data = self.stock_predictor.preprocess_data(
                self.stock_predictor.data, self.sequence_length, scaler_type
            )
            self.progress_updated.emit(15)
            
            self.message_updated.emit("构建模型...")
            self.stock_predictor.build_model(
                hidden_size=self.stock_predictor.hidden_size,
                num_layers=self.stock_predictor.num_layers,
                use_attention=self.use_attention
            )
            self.progress_updated.emit(25)
            
            self.message_updated.emit("开始训练模型...")
            self.stock_predictor.train_model(
                X_train, y_train, self.epochs, self.batch_size, self.learning_rate
            )
            self.progress_updated.emit(70)
            
            self.message_updated.emit("评估模型性能...")
            metrics = self.stock_predictor.evaluate_model(X_test, y_test)
            self.progress_updated.emit(80)
            
            self.message_updated.emit("预测未来走势...")
            future_predictions = self.stock_predictor.predict_future(
                self.stock_predictor.data, self.sequence_length
            )
            self.progress_updated.emit(90)
            
            self.message_updated.emit("生成历史预测...")
            history_predictions = self.stock_predictor.predict_history(X_test, scaled_data)
            self.progress_updated.emit(95)
            
            self.message_updated.emit("生成图表...")
            self.progress_updated.emit(100)
            
            self.prediction_completed.emit(metrics, future_predictions, history_predictions, self.stock_predictor.data)
        except Exception as e:
            self.message_updated.emit(f"错误: {str(e)}")

class MplCanvas(FigureCanvas):
    """Matplotlib画布，增强缩放和交互功能"""
    def __init__(self, parent=None, width=5, height=4, dpi=100):
        self.fig = Figure(figsize=(width, height), dpi=dpi)
        self.axes = self.fig.add_subplot(111)
        super(MplCanvas, self).__init__(self.fig)
        self.fig.tight_layout()
        
        # 初始化缩放状态
        self.original_limits = None
        
        # 连接鼠标滚轮事件
        self.mpl_connect('scroll_event', self.on_scroll)
        
        # 使用Matplotlib内置的导航工具栏
        self.toolbar = NavigationToolbar2QT(self, parent)
        parent.layout().addWidget(self.toolbar)  # 将工具栏添加到父布局
    
    def setup_zoom_controls(self):
        """设置缩放控制组件"""
        # 创建工具栏
        self.toolbar = self.fig.add_axes([0.1, 0.01, 0.8, 0.03])  # [left, bottom, width, height]
        
        # 放大按钮
        self.zoom_in = Button(self.toolbar, '放大')
        self.zoom_in.on_clicked(self.zoom_in_func)
        
        # 缩小按钮
        self.zoom_out = Button(self.toolbar, '缩小')
        self.zoom_out.on_clicked(self.zoom_out_func)
        
        # 重置按钮
        self.reset = Button(self.toolbar, '重置')
        self.reset.on_clicked(self.reset_zoom)
        
        # 保存原始坐标范围
        self.original_limits = (self.axes.get_xlim(), self.axes.get_ylim())
    
    def zoom_in_func(self, event):
        """放大图表"""
        self._zoom_chart(scale_factor=0.8)
    
    def zoom_out_func(self, event):
        """缩小图表"""
        self._zoom_chart(scale_factor=1.2)
    
    def reset_zoom(self, event):
        """重置图表缩放"""
        if self.original_limits:
            self.axes.set_xlim(self.original_limits[0])
            self.axes.set_ylim(self.original_limits[1])
            self.draw() 
    
    def on_scroll(self, event):
        """处理鼠标滚轮事件"""
        if event.inaxes != self.axes:
            return
            
        # 获取鼠标位置和当前坐标轴范围
        x, y = event.xdata, event.ydata
        xlim = self.axes.get_xlim()
        ylim = self.axes.get_ylim()
        
        # 根据滚轮方向确定缩放因子
        scale_factor = 0.8 if event.step > 0 else 1.2
        
        # 计算新的坐标范围，保持鼠标位置不变
        x_range = xlim[1] - xlim[0]
        y_range = ylim[1] - ylim[0]
        
        new_x_range = x_range * scale_factor
        new_y_range = y_range * scale_factor
        
        # 保持鼠标位置在缩放后的相对位置不变
        x_frac = (x - xlim[0]) / x_range
        y_frac = (y - ylim[0]) / y_range
        
        new_xlim = (x - new_x_range * x_frac, x + new_x_range * (1 - x_frac))
        new_ylim = (y - new_y_range * y_frac, y + new_y_range * (1 - y_frac))
        
        # 设置新的坐标范围并重绘
        self.axes.set_xlim(new_xlim)
        self.axes.set_ylim(new_ylim)
        self.draw()
    
    def _zoom_chart(self, scale_factor):
        """通用缩放函数"""
        xlim = self.axes.get_xlim()
        ylim = self.axes.get_ylim()
        
        x_range = xlim[1] - xlim[0]
        y_range = ylim[1] - ylim[0]
        
        x_mid = (xlim[0] + xlim[1]) / 2
        y_mid = (ylim[0] + ylim[1]) / 2
        
        new_x_range = x_range * scale_factor
        new_y_range = y_range * scale_factor
        
        new_xlim = (x_mid - new_x_range/2, x_mid + new_x_range/2)
        new_ylim = (y_mid - new_y_range/2, y_mid + new_y_range/2)
        
        self.axes.set_xlim(new_xlim)
        self.axes.set_ylim(new_ylim)
        self.draw()
    
    def setup_zoom_controls(self):
        """设置缩放控制组件"""
        # 创建工具栏
        self.toolbar = self.fig.add_axes([0.1, 0.01, 0.8, 0.03])  # [left, bottom, width, height]
        
        # 放大按钮
        self.zoom_in = Button(self.toolbar, '放大')
        self.zoom_in.on_clicked(self.zoom_in_func)
        
        # 缩小按钮
        self.zoom_out = Button(self.toolbar, '缩小')
        self.zoom_out.on_clicked(self.zoom_out_func)
        
        # 重置按钮
        self.reset = Button(self.toolbar, '重置')
        self.reset.on_clicked(self.reset_zoom)
        
        # 保存原始坐标范围
        self.original_limits = (self.axes.get_xlim(), self.axes.get_ylim())
    
    def zoom_in_func(self, event):
        """放大图表"""
        xlim = self.axes.get_xlim()
        ylim = self.axes.get_ylim()
        
        x_range = xlim[1] - xlim[0]
        y_range = ylim[1] - ylim[0]
        
        # 缩小范围（放大）
        x_mid = (xlim[0] + xlim[1]) / 2
        y_mid = (ylim[0] + ylim[1]) / 2
        
        new_x_range = x_range * 0.8
        new_y_range = y_range * 0.8
        
        new_xlim = (x_mid - new_x_range/2, x_mid + new_x_range/2)
        new_ylim = (y_mid - new_y_range/2, y_mid + new_y_range/2)
        
        self.axes.set_xlim(new_xlim)
        self.axes.set_ylim(new_ylim)
        self.draw()
    
    def zoom_out_func(self, event):
        """缩小图表"""
        xlim = self.axes.get_xlim()
        ylim = self.axes.get_ylim()
        
        x_range = xlim[1] - xlim[0]
        y_range = ylim[1] - ylim[0]
        
        # 扩大范围（缩小）
        x_mid = (xlim[0] + xlim[1]) / 2
        y_mid = (ylim[0] + ylim[1]) / 2
        
        new_x_range = x_range * 1.2
        new_y_range = y_range * 1.2
        
        new_xlim = (x_mid - new_x_range/2, x_mid + new_x_range/2)
        new_ylim = (y_mid - new_y_range/2, y_mid + new_y_range/2)
        
        self.axes.set_xlim(new_xlim)
        self.axes.set_ylim(new_ylim)
        self.draw()
    
    def reset_zoom(self, event):
        """重置图表缩放"""
        if self.original_limits:
            self.axes.set_xlim(self.original_limits[0])
            self.axes.set_ylim(self.original_limits[1])
            self.draw()

class StockPredictor:
    """股票/指数预测器"""
    def __init__(self, index_code: str, prediction_days: int = 7):
        """
        初始化股票预测器
        
        Args:
            index_code: 指数代码，例如"sh000001"代表上证指数
            prediction_days: 预测的交易日天数
        """
        self.index_code = index_code
        self.prediction_days = prediction_days
        self.scaler = None
        self.model = None
        self.data = None
        self.index_name = {
            "sh000001": "上证指数",
            "sh000300": "沪深300",
            "sh000905": "中证500",
            "sz399001": "深证成指",
            "sz399006": "创业板指"
        }.get(index_code, index_code)
        
        # 交易日检查器
        self.trading_day_checker = AShareTradingDayChecker()
        
        # 默认模型参数
        self.hidden_size = 50
        self.num_layers = 2
        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        
    def fetch_data(self, start_date: str = "2010-01-01", end_date: Optional[str] = None, 
                  use_technical_indicators: bool = True) -> pd.DataFrame:
        """
        从AkShare获取指数数据，只保留交易日数据，并可选择添加技术指标
        
        Args:
            start_date: 开始日期，格式"YYYY-MM-DD"
            end_date: 结束日期，格式"YYYY-MM-DD"，默认为今天
            use_technical_indicators: 是否添加技术指标
            
        Returns:
            包含指数数据的DataFrame
        """
        if end_date is None:
            end_date = datetime.datetime.now().strftime("%Y-%m-%d")
            
        try:
            print(f"正在获取{self.index_name}({self.index_code})从{start_date}到{end_date}的交易日数据...")
            
            # 使用AkShare获取指数数据
            if self.index_code.startswith("sh"):  # 上交所指数
                index_data = ak.stock_zh_index_daily(symbol=self.index_code)
            elif self.index_code.startswith("sz"):  # 深交所指数
                index_data = ak.stock_zh_index_daily(symbol=self.index_code)
            else:
                raise ValueError(f"不支持的指数代码: {self.index_code}")
            
            # 确保日期格式正确
            index_data['date'] = pd.to_datetime(index_data['date'])
            
            # 筛选日期范围
            mask = (index_data['date'] >= start_date) & (index_data['date'] <= end_date)
            filtered_data = index_data.loc[mask]
            
            # 重命名列以匹配模型期望的格式
            filtered_data = filtered_data.rename(columns={
                'date': 'Date',
                'open': 'Open',
                'high': 'High',
                'low': 'Low',
                'close': 'Close',
                'volume': 'Volume'
            })
            
            # 设置日期为索引
            filtered_data.set_index('Date', inplace=True)
            
            # 只保留交易日数据
            filtered_data = self._filter_trading_days(filtered_data)
            
            # 添加技术指标
            if use_technical_indicators:
                filtered_data = self._add_technical_indicators(filtered_data)
            
            print(f"成功获取{len(filtered_data)}个交易日数据")
            
            if filtered_data.empty:
                raise ValueError(f"无法获取{self.index_name}的交易日数据，请检查指数代码是否正确")
                
            self.data = filtered_data
            return filtered_data
        except Exception as e:
            print(f"获取数据时出错: {e}")
            raise e
    
    def _filter_trading_days(self, data: pd.DataFrame) -> pd.DataFrame:
        """筛选出交易日数据"""
        trading_days = []
        for date in data.index:
            if self.trading_day_checker.is_trading_day(date.date()):
                trading_days.append(date)
        return data.loc[trading_days]
    
    def _add_technical_indicators(self, data: pd.DataFrame) -> pd.DataFrame:
        """添加常用技术指标"""
        df = data.copy()
        
        # 移动平均线
        df['MA5'] = df['Close'].rolling(window=5).mean()
        df['MA10'] = df['Close'].rolling(window=10).mean()
        df['MA20'] = df['Close'].rolling(window=20).mean()
        
        # MACD
        df['EMA12'] = df['Close'].ewm(span=12, adjust=False).mean()
        df['EMA26'] = df['Close'].ewm(span=26, adjust=False).mean()
        df['MACD'] = df['EMA12'] - df['EMA26']
        df['Signal'] = df['MACD'].ewm(span=9, adjust=False).mean()
        df['Histogram'] = df['MACD'] - df['Signal']
        
        # RSI (相对强弱指标)
        delta = df['Close'].diff()
        gain = delta.where(delta > 0, 0)
        loss = -delta.where(delta < 0, 0)
        avg_gain = gain.rolling(window=14).mean()
        avg_loss = loss.rolling(window=14).mean()
        rs = avg_gain / avg_loss
        df['RSI'] = 100 - (100 / (1 + rs))
        
        # Bollinger Bands
        df['BB_Middle'] = df['Close'].rolling(window=20).mean()
        df['BB_Upper'] = df['BB_Middle'] + 2 * df['Close'].rolling(window=20).std()
        df['BB_Lower'] = df['BB_Middle'] - 2 * df['Close'].rolling(window=20).std()
        
        # 移除包含NaN的行
        df.dropna(inplace=True)
        
        return df
    
    def preprocess_data(self, data: pd.DataFrame, sequence_length: int = 60, 
                        use_scaler: str = 'minmax') -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
        """
        预处理数据用于模型训练
        
        Args:
            data: 包含指数数据的DataFrame
            sequence_length: 时间序列长度
            use_scaler: 使用的缩放器类型 ('minmax' 或 'standard')
            
        Returns:
            X_train, y_train, X_test, y_test, scaler
        """
        # 选择特征 (默认使用收盘价)
        if 'Close' in data.columns:
            feature_data = data[['Close']].values
        else:
            feature_data = data.values
            
        # 初始化缩放器
        if use_scaler == 'minmax':
            self.scaler = MinMaxScaler(feature_range=(0, 1))
        else:
            self.scaler = StandardScaler()
            
        # 归一化处理
        scaled_data = self.scaler.fit_transform(feature_data)
        
        # 创建序列
        X, y = [], []
        for i in range(sequence_length, len(scaled_data)):
            X.append(scaled_data[i-sequence_length:i, 0])
            y.append(scaled_data[i, 0])
            
        X, y = np.array(X), np.array(y)
        
        # 划分训练集和测试集
        train_size = int(len(X) * 0.8)
        X_train, X_test = X[:train_size], X[train_size:]
        y_train, y_test = y[:train_size], y[train_size:]
        
        # 转换为PyTorch张量
        X_train = torch.FloatTensor(X_train).unsqueeze(2)  # [样本数, 序列长度, 特征数]
        y_train = torch.FloatTensor(y_train).unsqueeze(1)
        X_test = torch.FloatTensor(X_test).unsqueeze(2)
        y_test = torch.FloatTensor(y_test).unsqueeze(1)
        
        return X_train, y_train, X_test, y_test, scaled_data
    
    def build_model(self, hidden_size: int = 50, num_layers: int = 2, use_attention: bool = False) -> None:
        """
        构建LSTM模型
        
        Args:
            hidden_size: LSTM隐藏层大小
            num_layers: LSTM层数
            use_attention: 是否使用注意力机制
        """
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        
        self.model = LSTMModel(
            input_size=1, 
            hidden_size=hidden_size, 
            num_layers=num_layers, 
            output_size=1,
            use_attention=use_attention
        )
        
        self.model.to(self.device)
        
        # 定义损失函数和优化器
        self.criterion = nn.MSELoss()
        
    def train_model(self, X_train: torch.Tensor, y_train: torch.Tensor, epochs: int = 50, 
                   batch_size: int = 32, learning_rate: float = 0.001) -> None:
        """
        训练模型
        
        Args:
            X_train: 训练数据
            y_train: 训练标签
            epochs: 训练轮数
            batch_size: 批次大小
            learning_rate: 学习率
        """
        if self.model is None:
            raise ValueError("请先构建模型")
            
        print(f"开始训练模型，训练轮数: {epochs}, 批次大小: {batch_size}, 学习率: {learning_rate}")
        
        # 创建数据加载器
        dataset = TensorDataset(X_train, y_train)
        dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
        
        # 优化器
        self.optimizer = optim.Adam(self.model.parameters(), lr=learning_rate)
        
        # 学习率调度器
        scheduler = optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer, mode='min', factor=0.5, patience=5, 
        )
        
        self.model.train()
        for epoch in range(epochs):
            total_loss = 0
            for X_batch, y_batch in dataloader:
                X_batch, y_batch = X_batch.to(self.device), y_batch.to(self.device)
                
                # 前向传播
                self.optimizer.zero_grad()
                outputs = self.model(X_batch)
                loss = self.criterion(outputs, y_batch)
                
                # 反向传播和优化
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)
                self.optimizer.step()
                
                total_loss += loss.item()
            
            # 更新学习率
            scheduler.step(total_loss/len(dataloader))
            
            if (epoch+1) % 10 == 0:
                print(f'Epoch [{epoch+1}/{epochs}], Loss: {total_loss/len(dataloader):.6f}')
                
        print("模型训练完成")
    
    def evaluate_model(self, X_test: torch.Tensor, y_test: torch.Tensor) -> Dict[str, float]:
        """评估模型性能"""
        if self.model is None:
            raise ValueError("请先训练模型")
            
        self.model.eval()
        with torch.no_grad():
            X_test = X_test.to(self.device)
            predictions = self.model(X_test).cpu().numpy()
            
        # 反归一化
        y_test = y_test.cpu().numpy()
        predictions = self.scaler.inverse_transform(predictions)
        y_test = self.scaler.inverse_transform(y_test)
        
        # 计算指标
        mae = mean_absolute_error(y_test, predictions)
        mse = mean_squared_error(y_test, predictions)
        rmse = np.sqrt(mse)
        r2 = r2_score(y_test, predictions)
        
        # 确保所有指标都返回，避免键不存在的问题
        return {
            "MAE": mae,
            "MSE": mse,
            "RMSE": rmse,
            "r2": r2  # 统一使用小写 'r2'
        }
    
    def predict_history(self, X_test: torch.Tensor, scaled_data: np.ndarray) -> np.ndarray:
        """
        预测历史数据，用于可视化模型在训练数据上的表现
        
        Args:
            X_test: 测试数据
            scaled_data: 标准化后的数据
            
        Returns:
            历史预测值
        """
        if self.model is None:
            raise ValueError("请先训练模型")
            
        self.model.eval()
        with torch.no_grad():
            X_test = X_test.to(self.device)
            predictions = self.model(X_test).cpu().numpy()
            
        # 反归一化
        predictions = self.scaler.inverse_transform(predictions)
        
        # 创建完整的历史预测数组（包括训练集和测试集）
        history_predictions = np.empty_like(scaled_data)
        history_predictions[:, :] = np.nan
        
        # 填充测试集预测
        test_start_idx = len(scaled_data) - len(predictions)
        history_predictions[test_start_idx:, :] = predictions
        
        return history_predictions
    
    def predict_future(self, data: pd.DataFrame, sequence_length: int = 60) -> np.ndarray:
        """
        预测未来的指数走势，只生成交易日的预测
        
        Args:
            data: 包含指数数据的DataFrame
            sequence_length: 时间序列长度
            
        Returns:
            预测的未来指数价格（仅交易日）
        """
        if self.model is None:
            raise ValueError("请先训练模型")
            
        close_prices = data['Close'].values.reshape(-1, 1)
        scaled_data = self.scaler.transform(close_prices)
        
        # 获取最后一段序列用于预测
        last_sequence = scaled_data[-sequence_length:]
        last_sequence = torch.FloatTensor(last_sequence).unsqueeze(0).to(self.device)  # [1, 序列长度, 1]
        
        # 预测未来价格 - 只预测交易日
        future_predictions = []
        self.model.eval()
        
        # 获取最后一个交易日
        last_date = data.index[-1].date()
        current_date = last_date
        
        with torch.no_grad():
            current_sequence = last_sequence.clone()
            
            # 生成预测，直到达到所需的交易日数量
            while len(future_predictions) < self.prediction_days:
                current_date += relativedelta(days=1)
                
                # 只在交易日进行预测
                if self.trading_day_checker.is_trading_day(current_date):
                    next_day_prediction = self.model(current_sequence)
                    future_predictions.append(next_day_prediction.item())
                    
                    # 更新序列，将预测值添加到序列末尾，并移除序列开头的值
                    new_sequence = torch.roll(current_sequence, -1, dims=1)
                    new_sequence[0, -1, 0] = next_day_prediction
                    current_sequence = new_sequence
            
        # 反归一化
        future_predictions = np.array(future_predictions).reshape(-1, 1)
        future_predictions = self.scaler.inverse_transform(future_predictions)
        
        return future_predictions
    
    def save_model(self, path: str) -> None:
        """保存模型"""
        if self.model is None:
            raise ValueError("没有训练好的模型可供保存")
            
        model_state = {
            'model_state_dict': self.model.state_dict(),
            'scaler': self.scaler,
            'index_code': self.index_code,
            'hidden_size': self.hidden_size,
            'num_layers': self.num_layers,
            'prediction_days': self.prediction_days
        }
        
        torch.save(model_state, path)
        print(f"模型已保存至 {path}")
    
    def load_model(self, path: str) -> None:
        """加载模型"""
        if not os.path.exists(path):
            raise FileNotFoundError(f"模型文件不存在: {path}")
            
        model_state = torch.load(path, map_location=self.device)
        
        self.index_code = model_state['index_code']
        self.hidden_size = model_state['hidden_size']
        self.num_layers = model_state['num_layers']
        self.prediction_days = model_state['prediction_days']
        self.scaler = model_state['scaler']
        
        # 构建模型
        self.build_model(
            hidden_size=self.hidden_size,
            num_layers=self.num_layers
        )
        
        # 加载模型参数
        self.model.load_state_dict(model_state['model_state_dict'])
        self.model.eval()
        
        print(f"模型已从 {path} 加载")

class MainWindow(QMainWindow):
    def __init__(self):
        super().__init__()
        self.setWindowTitle("A股指数预测系统 (含交易日判断)")
        self.setGeometry(100, 100, 1200, 800)
        
        # 创建中心部件和布局
        self.central_widget = QWidget()
        self.setCentralWidget(self.central_widget)
        self.main_layout = QVBoxLayout(self.central_widget)
        
        # 创建选项卡控件
        self.tab_widget = QTabWidget()
        self.main_layout.addWidget(self.tab_widget)
        
        # 创建设置选项卡
        self.setup_tab = QWidget()
        self.setup_layout = QVBoxLayout(self.setup_tab)
        self.tab_widget.addTab(self.setup_tab, "预测设置")
        
        # 创建结果选项卡
        self.result_tab = QWidget()
        self.result_layout = QVBoxLayout(self.result_tab)  # 初始化布局
        self.tab_widget.addTab(self.result_tab, "预测结果")
        
        # 创建评估选项卡
        self.evaluation_tab = QWidget()
        self.evaluation_layout = QVBoxLayout(self.evaluation_tab)  # 初始化布局
        self.tab_widget.addTab(self.evaluation_tab, "模型评估")
        
        # 初始化设置选项卡
        self.init_setup_tab()
        
        # 初始化结果选项卡（修正重复定义）
        self.init_result_tab()
        
        # 初始化评估选项卡（修正重复定义）
        self.init_evaluation_tab()
        
        # 初始化预测器
        self.stock_predictor = None
        self.prediction_thread = None
        self.tuning_thread = None
        self.future_predictions = None
        self.history_predictions = None
    
    def init_result_tab(self):
        """初始化结果选项卡（修正布局初始化）"""
        # 清空原有布局（避免重复添加）
        for i in reversed(range(self.result_layout.count())):
            widget = self.result_layout.itemAt(i).widget()
            if widget:
                widget.setParent(None)
        
        # 创建结果显示布局
        result_display_layout = QVBoxLayout()
        
        # 创建图表显示区域
        self.chart_canvas = MplCanvas(self, width=8, height=5, dpi=100)
        result_display_layout.addWidget(self.chart_canvas)
        
        # 创建结果信息区域
        info_layout = QHBoxLayout()
        
        # 模型评估结果
        metrics_label = QLabel("模型评估结果:")
        self.metrics_text = QTextEdit()
        self.metrics_text.setReadOnly(True)
        info_layout.addWidget(metrics_label)
        info_layout.addWidget(self.metrics_text)
        
        # 预测结果表格
        prediction_label = QLabel("未来预测结果:")
        self.prediction_text = QTextEdit()
        self.prediction_text.setReadOnly(True)
        info_layout.addWidget(prediction_label)
        info_layout.addWidget(self.prediction_text)
        
        result_display_layout.addLayout(info_layout)
        result_display_layout.setStretch(0, 3)  # 图表占3份
        result_display_layout.setStretch(1, 1)  # 信息占1份
        
        # 添加到结果选项卡布局
        self.result_layout.addLayout(result_display_layout)
        
        # 按钮布局
        button_layout = QHBoxLayout()
        
        # 保存结果按钮
        self.save_result_button = QPushButton("保存结果")
        self.save_result_button.clicked.connect(self.save_results)
        button_layout.addWidget(self.save_result_button)
        
        # 重新绘制按钮
        self.redraw_button = QPushButton("重新绘制")
        self.redraw_button.clicked.connect(self.redraw_chart)
        button_layout.addWidget(self.redraw_button)
        
        self.result_layout.addLayout(button_layout)
        
    def init_setup_tab(self):
        """初始化设置选项卡"""
        # 创建参数设置布局
        param_layout = QGridLayout()
        
        # 指数选择
        index_label = QLabel("选择指数:")
        self.index_combo = QComboBox()
        self.index_combo.addItems([
            "上证指数 (sh000001)",
            "沪深300 (sh000300)",
            "中证500 (sh000905)",
            "深证成指 (sz399001)",
            "创业板指 (sz399006)"
        ])
        param_layout.addWidget(index_label, 0, 0)
        param_layout.addWidget(self.index_combo, 0, 1)
        
        # 预测天数
        days_label = QLabel("预测交易日数:")
        self.days_spin = QSpinBox()
        self.days_spin.setRange(1, 30)
        self.days_spin.setValue(7)
        param_layout.addWidget(days_label, 0, 2)
        param_layout.addWidget(self.days_spin, 0, 3)
        
        # 训练轮数
        epochs_label = QLabel("训练轮数 (Epochs):")
        self.epochs_spin = QSpinBox()
        self.epochs_spin.setRange(1, 200)
        self.epochs_spin.setValue(50)
        param_layout.addWidget(epochs_label, 1, 0)
        param_layout.addWidget(self.epochs_spin, 1, 1)
        
        # 批次大小
        batch_label = QLabel("批次大小 (Batch Size):")
        self.batch_spin = QSpinBox()
        self.batch_spin.setRange(1, 128)
        self.batch_spin.setValue(32)
        param_layout.addWidget(batch_label, 1, 2)
        param_layout.addWidget(self.batch_spin, 1, 3)
        
        # 序列长度
        seq_label = QLabel("序列长度:")
        self.seq_spin = QSpinBox()
        self.seq_spin.setRange(10, 100)
        self.seq_spin.setValue(60)
        param_layout.addWidget(seq_label, 2, 0)
        param_layout.addWidget(self.seq_spin, 2, 1)
        
        # 学习率
        lr_label = QLabel("学习率:")
        self.lr_spin = QDoubleSpinBox()
        self.lr_spin.setRange(0.0001, 0.01)
        self.lr_spin.setSingleStep(0.0001)
        self.lr_spin.setValue(0.001)
        self.lr_spin.setDecimals(4)
        param_layout.addWidget(lr_label, 2, 2)
        param_layout.addWidget(self.lr_spin, 2, 3)
        
        # 模型参数组
        model_group = QGroupBox("模型参数")
        model_layout = QGridLayout()
        
        # 隐藏层大小
        hidden_label = QLabel("隐藏层大小:")
        self.hidden_spin = QSpinBox()
        self.hidden_spin.setRange(10, 200)
        self.hidden_spin.setValue(50)
        model_layout.addWidget(hidden_label, 0, 0)
        model_layout.addWidget(self.hidden_spin, 0, 1)
        
        # LSTM层数
        layers_label = QLabel("LSTM层数:")
        self.layers_spin = QSpinBox()
        self.layers_spin.setRange(1, 5)
        self.layers_spin.setValue(2)
        model_layout.addWidget(layers_label, 0, 2)
        model_layout.addWidget(self.layers_spin, 0, 3)
        
        # 使用注意力机制
        self.attention_checkbox = QCheckBox("使用注意力机制")
        model_layout.addWidget(self.attention_checkbox, 1, 0, 1, 2)
        
        # 使用技术指标
        self.tech_indicators_checkbox = QCheckBox("使用技术指标")
        self.tech_indicators_checkbox.setChecked(True)
        model_layout.addWidget(self.tech_indicators_checkbox, 1, 2, 1, 2)
        
        # 使用标准化而非归一化
        self.standard_scaler_checkbox = QCheckBox("使用标准化而非归一化")
        model_layout.addWidget(self.standard_scaler_checkbox, 2, 0, 1, 4)
        
        model_group.setLayout(model_layout)
        param_layout.addWidget(model_group, 3, 0, 1, 4)
        
        # 添加参数布局到设置选项卡
        self.setup_layout.addLayout(param_layout)
        
        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setValue(0)
        self.setup_layout.addWidget(self.progress_bar)
        
        # 日志显示
        self.log_text = QTextEdit()
        self.log_text.setReadOnly(True)
        self.setup_layout.addWidget(self.log_text)
        
        # 按钮布局
        button_layout = QHBoxLayout()
        
        # 预测按钮
        self.predict_button = QPushButton("开始预测")
        self.predict_button.clicked.connect(self.start_prediction)
        button_layout.addWidget(self.predict_button)
        
        # 超参数调优按钮
        self.tune_button = QPushButton("超参数调优")
        self.tune_button.clicked.connect(self.start_hyperparameter_tuning)
        button_layout.addWidget(self.tune_button)
        
        # 加载模型按钮
        self.load_model_button = QPushButton("加载模型")
        self.load_model_button.clicked.connect(self.load_model)
        button_layout.addWidget(self.load_model_button)
        
        # 保存模型按钮
        self.save_model_button = QPushButton("保存模型")
        self.save_model_button.setEnabled(False)
        button_layout.addWidget(self.save_model_button)
        
        self.setup_layout.addLayout(button_layout)
        
        # 添加间距
        self.setup_layout.addStretch(1)
        
    def init_evaluation_tab(self):
        """初始化评估选项卡（修正布局初始化）"""
        # 清空原有布局（避免重复添加）
        for i in reversed(range(self.evaluation_layout.count())):
            widget = self.evaluation_layout.itemAt(i).widget()
            if widget:
                widget.setParent(None)
        
        # 创建评估显示布局
        evaluation_layout = QHBoxLayout()
        
        # 创建图表显示区域
        self.eval_canvas = MplCanvas(self, width=8, height=5, dpi=100)
        evaluation_layout.addWidget(self.eval_canvas)
        
        # 创建评估信息区域
        eval_info_layout = QVBoxLayout()
        
        # 误差分布
        error_label = QLabel("误差分布:")
        self.error_text = QTextEdit()
        self.error_text.setReadOnly(True)
        eval_info_layout.addWidget(error_label)
        eval_info_layout.addWidget(self.error_text)
        
        # 评估指标
        indicators_label = QLabel("评估指标:")
        self.indicators_text = QTextEdit()
        self.indicators_text.setReadOnly(True)
        eval_info_layout.addWidget(indicators_label)
        eval_info_layout.addWidget(self.indicators_text)
        
        evaluation_layout.addLayout(eval_info_layout)
        
        # 添加到评估选项卡布局
        self.evaluation_layout.addLayout(evaluation_layout)
        
    def start_prediction(self):
        """开始预测过程"""
        # 禁用预测按钮
        self.predict_button.setEnabled(False)
        self.tune_button.setEnabled(False)
        self.load_model_button.setEnabled(False)
        self.log_text.clear()
        
        # 获取用户选择
        index_text = self.index_combo.currentText()
        index_code = index_text.split(" ")[1][1:-1]  # 提取括号中的代码
        prediction_days = self.days_spin.value()
        epochs = self.epochs_spin.value()
        batch_size = self.batch_spin.value()
        sequence_length = self.seq_spin.value()
        learning_rate = self.lr_spin.value()
        hidden_size = self.hidden_spin.value()
        num_layers = self.layers_spin.value()
        use_attention = self.attention_checkbox.isChecked()
        use_technical_indicators = self.tech_indicators_checkbox.isChecked()
        use_standard_scaler = self.standard_scaler_checkbox.isChecked()
        
        # 创建预测器
        self.stock_predictor = StockPredictor(index_code, prediction_days)
        self.stock_predictor.hidden_size = hidden_size
        self.stock_predictor.num_layers = num_layers
        
        # 记录日志
        self.log_text.append(f"开始预测 {self.stock_predictor.index_name}...")
        self.log_text.append(f"预测交易日数: {prediction_days}")
        self.log_text.append(f"训练轮数: {epochs}")
        self.log_text.append(f"批次大小: {batch_size}")
        self.log_text.append(f"序列长度: {sequence_length}")
        self.log_text.append(f"学习率: {learning_rate}")
        self.log_text.append(f"隐藏层大小: {hidden_size}")
        self.log_text.append(f"LSTM层数: {num_layers}")
        self.log_text.append(f"使用注意力机制: {use_attention}")
        self.log_text.append(f"使用技术指标: {use_technical_indicators}")
        self.log_text.append(f"使用标准化: {use_standard_scaler}")
        
        # 创建预测线程
        self.prediction_thread = PredictionThread(
            self.stock_predictor, epochs, batch_size, sequence_length, 
            learning_rate, use_attention, use_technical_indicators, use_standard_scaler
        )
        
        # 连接信号
        self.prediction_thread.progress_updated.connect(self.update_progress)
        self.prediction_thread.message_updated.connect(self.update_log)
        self.prediction_thread.prediction_completed.connect(self.show_results)
        
        # 启动线程
        self.prediction_thread.start()
        
    def start_hyperparameter_tuning(self):
        """开始超参数调优"""
        # 禁用按钮
        self.predict_button.setEnabled(False)
        self.tune_button.setEnabled(False)
        self.load_model_button.setEnabled(False)
        self.log_text.clear()
        
        # 获取用户选择
        index_text = self.index_combo.currentText()
        index_code = index_text.split(" ")[1][1:-1]  # 提取括号中的代码
        prediction_days = self.days_spin.value()
        sequence_length = self.seq_spin.value()
        use_attention = self.attention_checkbox.isChecked()
        use_technical_indicators = self.tech_indicators_checkbox.isChecked()
        use_standard_scaler = self.standard_scaler_checkbox.isChecked()
        
        # 创建预测器
        self.stock_predictor = StockPredictor(index_code, prediction_days)
        
        # 记录日志
        self.log_text.append(f"开始{self.stock_predictor.index_name}的超参数调优...")
        
        # 获取数据
        try:
            self.stock_predictor.fetch_data(use_technical_indicators=use_technical_indicators)
            self.log_text.append(f"成功获取{len(self.stock_predictor.data)}个交易日历史数据")
            
            # 预处理数据
            scaler_type = 'standard' if use_standard_scaler else 'minmax'
            X_train, y_train, X_test, y_test, _ = self.stock_predictor.preprocess_data(
                self.stock_predictor.data, sequence_length, scaler_type
            )
            self.log_text.append("数据预处理完成")
            
            # 创建调优线程
            self.tuning_thread = HyperparameterTuningThread(
                self.stock_predictor, X_train, y_train, X_test, y_test
            )
            
            # 连接信号
            self.tuning_thread.progress_updated.connect(self.update_progress)
            self.tuning_thread.message_updated.connect(self.update_log)
            self.tuning_thread.tuning_completed.connect(self.show_tuning_results)
            
            # 启动线程
            self.tuning_thread.start()
            
        except Exception as e:
            self.log_text.append(f"获取数据或预处理时出错: {str(e)}")
            self.predict_button.setEnabled(True)
            self.tune_button.setEnabled(True)
            self.load_model_button.setEnabled(True)
    
    def show_tuning_results(self, results):
        """显示超参数调优结果"""
        self.log_text.append("超参数调优结果:")
        best_params = results['best_params']
        self.log_text.append(f"最佳参数: {best_params}")
        
        # 更新界面参数
        self.hidden_spin.setValue(best_params['hidden_size'])
        self.layers_spin.setValue(best_params['num_layers'])
        self.epochs_spin.setValue(best_params['epochs'])
        self.batch_spin.setValue(best_params['batch_size'])
        self.lr_spin.setValue(best_params['learning_rate'])
        
        self.log_text.append("超参数调优完成，已更新为最佳参数")
        self.progress_bar.setValue(100)
        
        # 启用按钮
        self.predict_button.setEnabled(True)
        self.tune_button.setEnabled(True)
        self.load_model_button.setEnabled(True)
        
        # 显示评估选项卡
        self.tab_widget.setCurrentIndex(2)
        self.show_evaluation(results)
    
    def show_evaluation(self, results):
        """显示评估结果"""
        # 绘制超参数调优结果图表
        self.eval_canvas.axes.clear()
        
        # 提取数据
        params = results['results']
        rmse_values = [p['RMSE'] for p in params]
        param_names = [f"隐藏层:{p['hidden_size']}, 层数:{p['num_layers']}" for p in params]
        
        # 绘制RMSE柱状图
        self.eval_canvas.axes.bar(param_names, rmse_values, color='skyblue')
        self.eval_canvas.axes.set_title('不同超参数组合的RMSE对比')
        self.eval_canvas.axes.set_xlabel('超参数组合')
        self.eval_canvas.axes.set_ylabel('RMSE')
        self.eval_canvas.axes.tick_params(axis='x', rotation=45)
        self.eval_canvas.fig.tight_layout()
        self.eval_canvas.draw()
        
        # 显示评估指标
        indicators_text = "超参数调优评估指标:\n\n"
        indicators_text += f"最佳 RMSE: {min(rmse_values):.4f}\n"
        indicators_text += f"平均 RMSE: {sum(rmse_values)/len(rmse_values):.4f}\n"
        indicators_text += f"最佳参数: {results['best_params']}"
        self.indicators_text.setText(indicators_text)
    
    def show_results(self, metrics, future_predictions, history_predictions, data):
        """显示预测结果"""
        # 保存结果
        self.future_predictions = future_predictions
        self.history_predictions = history_predictions
        
        # 更新日志
        self.log_text.append("预测完成！")
        
        # 启用按钮
        self.predict_button.setEnabled(True)
        self.tune_button.setEnabled(True)
        self.load_model_button.setEnabled(True)
        self.save_model_button.setEnabled(True)
        
        # 更新进度条
        self.progress_bar.setValue(100)
        
        # 显示结果选项卡
        self.tab_widget.setCurrentIndex(1)
        
        # 显示评估指标 - 增强防御性
        metrics_text = f"模型评估指标:\n\n"
        metrics_text += f"平均绝对误差 (MAE): {metrics.get('MAE', 'N/A'):.4f}\n"
        metrics_text += f"均方误差 (MSE): {metrics.get('MSE', 'N/A'):.4f}\n"
        metrics_text += f"均方根误差 (RMSE): {metrics.get('RMSE', 'N/A'):.4f}\n"
        
        # 检查R²指标是否存在
        if 'r2' in metrics:
            metrics_text += f"决定系数 (R²): {metrics['r2']:.4f}"
        elif 'R2' in metrics:
            metrics_text += f"决定系数 (R²): {metrics['R2']:.4f}"
        else:
            metrics_text += f"决定系数 (R²): N/A"
        
        self.metrics_text.setText(metrics_text)
        
        # 显示未来预测结果
        prediction_text = f"{self.stock_predictor.index_name}未来{self.stock_predictor.prediction_days}个交易日预测结果:\n\n"
        last_date = data.index[-1]
        
        for i, price in enumerate(future_predictions):
            # 计算预测日期（只考虑交易日）
            trading_day_checker = AShareTradingDayChecker()
            prediction_date = last_date
            days_added = 0
            
            while days_added <= i:
                prediction_date += relativedelta(days=1)
                if trading_day_checker.is_trading_day(prediction_date.date()):
                    days_added += 1
            
            prediction_text += f"{prediction_date.date()}: {price[0]:.2f}\n"
            
        self.prediction_text.setText(prediction_text)
        
        # 绘制图表
        self.redraw_chart()
        
    def redraw_chart(self):
        """重新绘制图表"""
        if self.future_predictions is None or self.history_predictions is None:
            QMessageBox.warning(self, "警告", "没有预测结果可供显示")
            return
            
        # 清除画布
        self.chart_canvas.axes.clear()
        
        # 获取原始数据
        close_prices = self.stock_predictor.data['Close'].values
        
        # 创建日期范围
        dates = self.stock_predictor.data.index
        future_dates = []
        
        # 生成未来交易日日期
        last_date = dates[-1]
        trading_day_checker = AShareTradingDayChecker()
        current_date = last_date
        
        for _ in range(len(self.future_predictions)):
            while True:
                current_date += relativedelta(days=1)
                if trading_day_checker.is_trading_day(current_date.date()):
                    future_dates.append(current_date)
                    break
        
        # 绘制原始价格
        self.chart_canvas.axes.plot(dates, close_prices, label='实际价格', color='blue')
        
        # 绘制历史预测
        valid_history = self.history_predictions[~np.isnan(self.history_predictions)]
        valid_dates = dates[-len(valid_history):]
        self.chart_canvas.axes.plot(valid_dates, valid_history, label='历史预测', color='green', linestyle='--')
        
        # 绘制未来预测
        self.chart_canvas.axes.plot(future_dates, self.future_predictions, label='未来预测', color='red', linestyle='--')
        
        # 添加标题和标签
        self.chart_canvas.axes.set_title(f"{self.stock_predictor.index_name}价格预测 (交易日)")
        self.chart_canvas.axes.set_xlabel('日期')
        self.chart_canvas.axes.set_ylabel('价格')
        self.chart_canvas.axes.legend()
        self.chart_canvas.axes.grid(True)
        
        # 自动调整x轴日期显示
        self.chart_canvas.fig.autofmt_xdate()
        self.chart_canvas.fig.tight_layout()
        self.chart_canvas.draw()
        
    def save_results(self):
        """保存预测结果"""
        if self.future_predictions is None:
            QMessageBox.warning(self, "警告", "没有预测结果可供保存")
            return
            
        # 打开文件对话框
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存预测结果", "", "CSV文件 (*.csv);;Excel文件 (*.xlsx)"
        )
        
        if not file_path:
            return
            
        try:
            # 获取原始数据
            close_prices = self.stock_predictor.data['Close']
            
            # 创建未来交易日日期
            last_date = self.stock_predictor.data.index[-1]
            trading_day_checker = AShareTradingDayChecker()
            future_dates = []
            current_date = last_date
            
            for _ in range(len(self.future_predictions)):
                while True:
                    current_date += relativedelta(days=1)
                    if trading_day_checker.is_trading_day(current_date.date()):
                        future_dates.append(current_date)
                        break
            
            # 创建预测数据DataFrame
            future_data = pd.DataFrame({
                '日期': future_dates,
                '预测价格': [price[0] for price in self.future_predictions]
            })
            
            # 创建历史预测DataFrame
            history_data = pd.DataFrame({
                '日期': self.stock_predictor.data.index,
                '实际价格': close_prices,
                '历史预测价格': self.history_predictions.flatten()
            })
            
            # 保存到文件
            if file_path.endswith('.csv'):
                future_data.to_csv(file_path, index=False, encoding='utf-8-sig')
            elif file_path.endswith('.xlsx'):
                with pd.ExcelWriter(file_path) as writer:
                    future_data.to_excel(writer, sheet_name='未来预测', index=False)
                    history_data.to_excel(writer, sheet_name='历史预测', index=False)
                    
            self.log_text.append(f"预测结果已保存至 {file_path}")
            QMessageBox.information(self, "成功", f"预测结果已保存至 {file_path}")
        except Exception as e:
            self.log_text.append(f"保存结果时出错: {str(e)}")
            QMessageBox.critical(self, "错误", f"保存结果时出错: {str(e)}")
            
    def load_model(self):
        """加载预训练模型"""
        # 打开文件对话框
        file_path, _ = QFileDialog.getOpenFileName(
            self, "加载模型", "", "模型文件 (*.pth;*.pt)"
        )
        
        if not file_path:
            return
            
        try:
            # 禁用按钮
            self.predict_button.setEnabled(False)
            self.tune_button.setEnabled(False)
            self.load_model_button.setEnabled(False)
            self.log_text.clear()
            
            # 创建新的预测器
            self.stock_predictor = StockPredictor("sh000001")  # 默认上证指数
            
            # 加载模型
            self.stock_predictor.load_model(file_path)
            self.log_text.append(f"成功从 {file_path} 加载模型")
            
            # 更新界面显示
            index_names = {
                "sh000001": "上证指数 (sh000001)",
                "sh000300": "沪深300 (sh000300)",
                "sh000905": "中证500 (sh000905)",
                "sz399001": "深证成指 (sz399001)",
                "sz399006": "创业板指 (sz399006)"
            }
            
            # 更新指数选择
            index_text = index_names.get(self.stock_predictor.index_code, self.stock_predictor.index_code)
            index_idx = self.index_combo.findText(index_text)
            if index_idx >= 0:
                self.index_combo.setCurrentIndex(index_idx)
                
            # 更新预测天数
            self.days_spin.setValue(self.stock_predictor.prediction_days)
            
            # 更新模型参数
            self.hidden_spin.setValue(self.stock_predictor.hidden_size)
            self.layers_spin.setValue(self.stock_predictor.num_layers)
            
            # 获取数据
            self.stock_predictor.fetch_data()
            self.log_text.append(f"成功获取{len(self.stock_predictor.data)}个交易日历史数据")
            
            # 预测
            self.log_text.append("使用加载的模型进行预测...")
            sequence_length = self.seq_spin.value()
            _, _, X_test, _, scaled_data = self.stock_predictor.preprocess_data(
                self.stock_predictor.data, sequence_length
            )
            
            # 评估模型
            metrics = self.stock_predictor.evaluate_model(X_test, X_test)  # 使用测试数据评估
            
            # 预测未来和历史
            future_predictions = self.stock_predictor.predict_future(
                self.stock_predictor.data, sequence_length
            )
            history_predictions = self.stock_predictor.predict_history(X_test, scaled_data)
            
            # 显示结果
            self.show_results(metrics, future_predictions, history_predictions, self.stock_predictor.data)
            
        except Exception as e:
            self.log_text.append(f"加载模型时出错: {str(e)}")
            QMessageBox.critical(self, "错误", f"加载模型时出错: {str(e)}")
            self.predict_button.setEnabled(True)
            self.tune_button.setEnabled(True)
            self.load_model_button.setEnabled(True)
            
    def save_model(self):
        """保存当前模型"""
        if self.stock_predictor is None or self.stock_predictor.model is None:
            QMessageBox.warning(self, "警告", "没有训练好的模型可供保存")
            return
            
        # 打开文件对话框
        file_path, _ = QFileDialog.getSaveFileName(
            self, "保存模型", "", "模型文件 (*.pth)"
        )
        
        if not file_path:
            return
            
        try:
            # 保存模型
            self.stock_predictor.save_model(file_path)
            self.log_text.append(f"模型已保存至 {file_path}")
            QMessageBox.information(self, "成功", f"模型已保存至 {file_path}")
        except Exception as e:
            self.log_text.append(f"保存模型时出错: {str(e)}")
            QMessageBox.critical(self, "错误", f"保存模型时出错: {str(e)}")
            
    def update_progress(self, value):
        """更新进度条"""
        self.progress_bar.setValue(value)
        
    def update_log(self, message):
        """更新日志文本"""
        self.log_text.append(message)

if __name__ == "__main__":
    app = QApplication(sys.argv)
    window = MainWindow()
    window.show()
    sys.exit(app.exec_())  